import os  # 导入操作系统相关模块，用于处理文件路径、环境变量等操作
import asyncio  # 导入异步操作模块，用于实现异步操作
from lightrag import LightRAG, QueryParam  # 从lightrag库导入LightRAG类和QueryParam类，用于初始化RAG系统和查询
from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed, \
    create_openai_async_client  # 导入OpenAI的GPT-4o-mini模型和嵌入函数
from lightrag.utils import EmbeddingFunc  # 导入嵌入函数封装类EmbeddingFunc，用于处理文本的嵌入
import numpy as np  # 导入numpy库，用于数组操作，特别是嵌入返回的数组
from lightrag.kg.shared_storage import initialize_pipeline_status  # 导入初始化管道状态的函数，用于设置状态
from dotenv import load_dotenv  # 新增：导入 dotenv 模块，用于加载 .env 文件中的环境变量

# 新增：加载 .env 文件中的环境变量
load_dotenv()
# 处理异步操作的模块（如果是在Jupyter Notebook环境中运行）
#########
# Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
# import nest_asyncio
# nest_asyncio.apply()
#########

WORKING_DIR = "./mongodb_test_dir"  # 设置工作目录，存储MongoDB相关的数据
if not os.path.exists(WORKING_DIR):  # 如果工作目录不存在
    os.mkdir(WORKING_DIR)  # 创建工作目录

# 设置环境变量，用于配置OpenAI API密钥和MongoDB连接信息
os.environ["OPENAI_API_KEY"] = "sk-"  # 设置OpenAI API密钥
os.environ["MONGO_URI"] = "mongodb://0.0.0.0:27017/?directConnection=true"  # 设置MongoDB连接URI
os.environ["MONGO_DATABASE"] = "LightRAG"  # 设置MongoDB数据库名称
os.environ["MONGO_KG_COLLECTION"] = "MDB_KG"  # 设置MongoDB中用于存储知识图谱的集合名称

# 嵌入模型的配置和函数
# EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large")  # 从环境变量中获取嵌入模型的名称，如果没有则使用默认值
EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192))  # 获取最大token数量，如果没有则使用默认值8192

# 定义一个异步嵌入函数，使用OpenAI的嵌入模型进行文本处理
async def embedding_func(texts: list[str]) -> np.ndarray:

    return await openai_embed(
        texts,
        model=os.getenv("EMBEDDING_MODEL"),  # 从环境变量中获取嵌入模型的名称
        api_key=os.getenv("EMBEDDING_BINDING_API_KEY"),  # 从环境变量中获取嵌入模型的 API 密钥
        base_url=os.getenv("EMBEDDING_BINDING_HOST"),  # 从环境变量中获取嵌入模型的 API 基础 URL

    )

# 定义一个异步函数，用于获取嵌入的维度
async def get_embedding_dimension():
    test_text = ["This is a test sentence."]  # 这是一个测试句子
    embedding = await embedding_func(test_text)  # 获取嵌入
    return embedding.shape[1]  # 返回嵌入的维度

# 定义一个异步函数，用于创建嵌入函数实例
async def create_embedding_function_instance():
    # 获取嵌入的维度
    embedding_dimension = await get_embedding_dimension()
    # 创建嵌入函数实例
    return EmbeddingFunc(
        embedding_dim=embedding_dimension,  # 设置嵌入维度
        max_token_size=EMBEDDING_MAX_TOKEN_SIZE,  # 设置最大token数量
        func=embedding_func,  # 设置实际的嵌入函数
    )

# 定义一个异步函数，用于初始化RAG（检索增强生成）系统
async def initialize_rag():
    embedding_func_instance = await create_embedding_function_instance()  # 创建嵌入函数实例

    # 初始化LightRAG实例
    rag = LightRAG(
        working_dir=WORKING_DIR,  # 设置工作目录
        llm_model_func=create_openai_async_client,  # 设置使用的LLM（语言生成模型），此处使用GPT-4o-mini
        embedding_func=embedding_func_instance,  # 设置嵌入函数实例
        graph_storage="MongoGraphStorage",  # 设置知识图谱存储方式，使用MongoDB存储
        log_level="DEBUG",  # 设置日志级别为DEBUG，用于调试
    )

    await rag.initialize_storages()  # 初始化存储
    await initialize_pipeline_status()  # 初始化管道状态
    return rag  # 返回初始化后的RAG实例

# 主函数
def main():
    # 初始化RAG实例
    rag = asyncio.run(initialize_rag())

    # 读取文本文件并插入到RAG系统中
    with open("./book.txt", "r", encoding="utf-8") as f:
        rag.insert(f.read())  # 将文件内容插入RAG系统

    # 执行不同模式的查询并打印结果

    # 执行朴素搜索
    print(
        rag.query(
            "这个故事的主要主题是什么？", param=QueryParam(mode="naive")  # 进行朴素查询
        )
    )

    # 执行局部搜索
    print(
        rag.query(
            "这个故事的主要主题是什么？", param=QueryParam(mode="local")  # 进行局部查询
        )
    )

    # 执行全局搜索
    print(
        rag.query(
            "这个故事的主要主题是什么？", param=QueryParam(mode="global")  # 进行全局查询
        )
    )

    # 执行混合搜索
    print(
        rag.query(
            "这个故事的主要主题是什么？", param=QueryParam(mode="hybrid")  # 进行混合查询
        )
    )

# 判断是否为主程序运行
if __name__ == "__main__":
    main()  # 运行主函数
